#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# glumpy is an OpenGL framework for the fast visualization of numpy arrays.
# Copyright (C) 2009-2012  Nicolas P. Rougier. All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 
# THIS SOFTWARE IS PROVIDED BY NICOLAS P. ROUGIER ''AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
# EVENT SHALL NICOLAS P. ROUGIER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# 
# The views and conclusions contained in the software and documentation are
# those of the authors and should not be interpreted as representing official
# policies, either expressed or implied, of Nicolas P. Rougier.
# -----------------------------------------------------------------------------
'''
'''
import numpy as np
import OpenGL.GL as gl
from scipy.spatial import cKDTree
from glumpy import figure, Trackball


if __name__ == '__main__':
    
    fig = figure(size=(800,800))
    trackball = Trackball(65, 135, 1., 2.0)

    @fig.event
    def on_mouse_drag(x, y, dx, dy, button):
        trackball.drag_to(x,y,dx,dy)
        fig.redraw()

    @fig.event
    def on_draw():
        fig.clear()
        trackball.push()
        gl.glEnable(gl.GL_POINT_SMOOTH)
        for i,P in enumerate([boids['position'],
                              boids['position_1'],
                              boids['position_2']]):
            gl.glColor(0+i*.25,0+i*.25,0+i*.25,0.5)
            gl.glPointSize(6.0-i)
            gl.glBegin(gl.GL_POINTS)
            for j in range(len(boids)):
                gl.glVertex( P[j,0], P[j,1], P[j,2] )
            gl.glEnd()
        trackball.pop()

    @fig.timer(60.0)
    def timer(dt):
        global t

        t += 0.5*dt
        target[...] = np.array([np.sin(t),np.sin(2*t),np.cos(3*t)])*.1

        t += 0.5*dt
        predator[...] = np.array([np.sin(t),np.sin(2*t),np.cos(3*t)])*.2

        boids['position_2'] = boids['position_1']
        boids['position_1'] = boids['position']
        n = len(boids)
        P = boids['position']
        V = boids['velocity']

        # Cohesion: steer to move toward the average position of local flockmates
        C = -(P - P.sum(axis=0)/n)

        # Alignment: steer towards the average heading of local flockmates
        A = -(V - V.sum(axis=0)/n)

        # Repulsion: steer to avoid crowding local flockmates
        D,I = cKDTree(P).query(P,5)
        M = np.repeat(D < 0.05, 3, axis=1).reshape(n,5,3)
        Z = np.repeat(P,5,axis=0).reshape(n,5,3)
        R = -((P[I]-Z)*M).sum(axis=1)

        # Target : Follow target
        T = target - P

        # Predator : Move away from predator
        dP = P - predator
        D = np.maximum(0, 0.1 - np.sqrt(dP[:,0]**2 +dP[:,1]**2+dP[:,2]**2) )
        D = np.repeat(D,3,axis=0).reshape(n,3)
        dP *= D

        boids['velocity'] += 0.0005*C + 0.01*A + 0.01*R + 0.0005*T + 0.0025*dP
        boids['position'] += boids['velocity']

        fig.redraw()


    t = 0
    n = 1000
    boids = np.zeros(n, [ ('position',   'f4', 3),
                          ('position_1', 'f4', 3),
                          ('position_2', 'f4', 3),
                          ('velocity',   'f4', 3)] )
    boids['position'] = np.random.uniform(-0.25, +0.25, (n,3))
    boids['velocity'] = np.random.uniform(-0.00, +0.00, (n,3))
    target   = np.zeros(3)
    predator = np.zeros(3)


    fig.show()
